from tracker_elastic_util import tracker_manager
from yolo_v3_util import yolo_detect

class track_and_detect:
    def __init__(self, yolo_path, siamrpn_path, target):
        self.frame_num = 0
        self.keyframe_num = 25
        self.is_keyframe = False
        self.tracker = tracker_manager(siamrpn_path, track_method='SiamRPN')
        self.detector = yolo_detect(yolo_path, target=target, draw=False)

    def updated(self, frame, gt=None):
        bbox = []
        if self.frame_num % self.keyframe_num == 0:
            bbox = self.detector.yolo_pipeline(frame)
            self.frame_num = 0
            self.is_keyframe = True
        else:
            self.is_keyframe = False
        self.frame_num += 1
        ret = self.tracker.tracker_update(frame, is_keyframe=self.is_keyframe, key_bbox=bbox, gt=gt)

    def get_iou_list(self):
        return self.tracker.iou_list